package xyz.zenheart.ffmpegrelay.config;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import javax.imageio.stream.FileImageOutputStream;
import java.io.File;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class WsHandler extends BinaryWebSocketHandler {

    static final Map<String, Map<String, WebSocketSession>> services = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("新的加入");
        String key = (String) session.getAttributes().get("key");
        if (Objects.isNull(key)) {
            return;
        }
        if (!services.containsKey(key)) {
            services.put(key, new ConcurrentHashMap<>());
        }
        services.get(key).put(session.getId(), session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        removeWS(session);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        removeWS(session);
    }

    private void removeWS(WebSocketSession session) {
        String key = (String) session.getAttributes().get("key");
        if (Objects.isNull(key)) {
            return;
        }
        services.get(key).remove(session.getId());
        if (services.get(key).isEmpty()) {
            services.get(key).remove(key);
        }
    }

    public void sendVideo(String user, byte[] data) {
        Map<String, WebSocketSession> map = services.get(user);
        if (Objects.isNull(map)) {
            return;
        }
        try {
            BinaryMessage binaryMessage = new BinaryMessage(data);
            for (Map.Entry<String, WebSocketSession> sessionEntry : map.entrySet()) {
                WebSocketSession session = sessionEntry.getValue();
                try {
                    if (session.isOpen()) {
                        session.sendMessage(binaryMessage);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void byte2image(byte[] data, String path) {
        if (data.length < 3 || path.equals("")) return;
        try {
            FileImageOutputStream imageOutput = new FileImageOutputStream(new File(path));
            imageOutput.write(data, 0, data.length);
            imageOutput.close();
            System.out.println("Make Picture success,Please find image in " + path);
        } catch (Exception ex) {
            System.out.println("Exception: " + ex);
            ex.printStackTrace();
        }
    }
}
